# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import datetime
import os

from absl import app, flags
from pymongo import MongoClient

import google.protobuf.json_format as json_format
from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog
from coin.support.proto_log.logic.util import run_from_og_log_archive

FLAGS = flags.FLAGS


class OnLog(object):
  def __init__(self, collection_name):
    self.num_records = 0
    db = MongoClient().og_log
    self._collection = getattr(db, collection_name)

  def on_log(self, timestamp, log):
    pb = OrderGatewayLog()
    pb.ParseFromString(log)
    data = json_format.MessageToDict(pb, preserving_proto_field_name=True)
    hr_ts = datetime.datetime.fromtimestamp(timestamp / 1e9).strftime('%Y%m%d %H:%M:%S.%f')
    data['record_timestamp'] = timestamp
    data['record_timestamp_human_readable'] = hr_ts
    self.num_records += 1
    self._collection.insert_one(data)


def launch(owner, market_type, exchange, trading_date, root_dir, machine, on_log):
  start_time = trading_date  # UTC 0 hour
  hours = 2400

  run_from_og_log_archive(on_log,
                          market_type,
                          exchange,
                          start_time,
                          start_time + datetime.timedelta(hours=hours),
                          root_dir,
                          machine,
                          owner=owner)
  return 0


def main(_):
  owner = FLAGS.owner
  market_type = FLAGS.market_type
  exchange = FLAGS.exchange
  trading_date = FLAGS.trading_date
  root_dir = os.path.expanduser(FLAGS.root_dir)
  machine = FLAGS.machine
  collection = FLAGS.collection
  assert market_type, '--market_type must be specified.'
  assert exchange, '--exchange must be specified.'
  assert trading_date, '--trading_date must be specified.'
  assert machine, '--machine must be specified.'
  assert collection, '--collection must be specified.'
  print('Running for %s %s ...' % (trading_date, machine))
  onlog = OnLog(collection)
  trading_date = datetime.datetime.strptime(trading_date, '%Y%m%d')
  launch(owner, market_type, exchange, trading_date, root_dir, machine, onlog.on_log)
  print(onlog.num_records)


if __name__ == '__main__':
  flags.DEFINE_string('owner', None, 'Account.')

  flags.DEFINE_string('market_type', None, 'Market type.')

  flags.DEFINE_string('exchange', None, 'Exchange name.')

  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('root_dir', '~/data/og_proto_log', 'root_dir.')

  flags.DEFINE_string('machine', None, 'Instance machine name.')

  flags.DEFINE_string('collection', None, 'collection name')

  app.run(main)
